"""
https://github.com/offu/WeRoBot/issues/486
腾讯的限制
[BUG] handle里面函数处理超过5秒会连续调用3次
"""

import werobot
import datetime
import sys
import os
import time
import redis
import pymongo as pm
import threading

sys.path.insert(0, '/home/yunpeng/code')
from PyCmpltrtok.common import sep, to_decode, to_int
from PyCmpltrtok.util_mongo import delete_many_by_user, get_history, enqueue
from common import MONGODB_NAME, VALUE
from fastapi_interface import text_gen

os.environ['R_HOST'] = '127.0.0.1'
os.environ['R_PORT'] = '6379'
os.environ['M_HOST'] = '127.0.0.1'
os.environ['M_PORT'] = '27017'
os.environ['M_USER'] = 'root'
os.environ['M_PWD'] = 'p1983mdA1-ei'
os.environ['W_PORT'] = '8081'

# 连接Redis
sep('Redis')
rdb = redis.Redis(os.environ['R_HOST'], int(os.environ['R_PORT']), 0, socket_timeout=3)
rdb.get('try_it')
sep('Redis OK')

# 连接Mongodb
sep('MongoDB')
mongo = pm.MongoClient(
    os.environ['M_HOST'], int(os.environ['M_PORT']), serverSelectionTimeoutMS=3000,
    username=os.environ['M_USER'], password=os.environ['M_PWD'], 
)
mdb = mongo[MONGODB_NAME]
get_history(mdb, 'user_xxxx', limit=1)  # try it
sep('MongoDB OK')

sep('Werobot')
robot = werobot.WeRoBot(token='kx8M7eODU4HiyAmw3a')
robot.config['APP_ID'] = 'wx47577a029b258f6b'
robot.config['ENCODING_AES_KEY'] = 'LjBrnwSvqGXx89Tszleo4sMYp8yMALcTArwhPjxoZGJ'
client = robot.client
sep('Werobot OK')

CHAR_LIMIT = 450
PROMPT_SHORT = '（请继续对话，或者输入逗号（,或，）来重置对话。）'
PROMPT = '（因平台限制，内容未完全展示，输入句号（.或。）来继续展示。请继续对话，或者输入逗号（,或，）来重置对话。）'
GO_ON_SET = set(['.', '。'])
RESET_SET = set([',', '，'])
WAIT_LAST_TURN = '（请等待上一轮对话完成。）'
THRESH_S = 20  # 单例锁最多等待这么多秒

if 0:
    """
    自定义菜单 需要 微信认证
    微信认证： 账号主体为个人，无法开通微信认证
    """
    client.create_menu({
        'button': [{
            'type': 'click',
            'name': '继续',
            'key': 'go_on'
        }]
    })


    @robot.key_click('go_on')
    def go_on(message):    
        pass
    

@robot.handler
def hello(message):
    
    sep()
    
    # 判断类型
    xtype = message.type
    print(datetime.datetime.now(), xtype, flush=True)
    if 'text' != xtype:
        return '目前仅支持文本信息！' + PROMPT_SHORT
    
    # 用户标识
    xuserid = message.source
    print('userid:', type(xuserid), xuserid, flush=True)
    
    # 输入文本
    xinput = message.content.strip()
    print('input:', xinput, flush=True)
    
    # 单例限制
    if xinput not in GO_ON_SET:
        xis_in = rdb.sismember('users_now_lock', xuserid)
        ts_now = time.time_ns()
        if xis_in:
            ts = rdb.hget('username2ts', xuserid)
            ts = to_int(ts)
            if ts and ts_now - ts < THRESH_S * 1e9:
                users_count = rdb.scard('users_now')
                xoutput = WAIT_LAST_TURN + PROMPT_SHORT + f'【活跃用户{users_count}人】'
                print('OUTPUT:', xoutput, flush=True)
                return xoutput
        rdb.hset('username2ts', xuserid, ts_now)
        rdb.sadd('users_now_lock', xuserid)
        
    def do_return(xoutput):
        """
        闭包函数，大部分返回走这里。
        """
        users_count = rdb.scard('users_now')
        xoutput += f'【活跃用户{users_count}人】'
        print('OUTPUT:', xoutput, flush=True)
        # 解除单例限制
        rdb.srem('users_now_lock', xuserid)
        return xoutput

    def get_result():
        """
        闭包函数，拿数据
        """
        # 最后结果
        xoutput = rdb.hget('username2output', xuserid)
        
        # 当前偏移
        xoffset = rdb.hget('username2offset', xuserid)
        xoffset = to_int(xoffset)
        xoffset_ = xoffset
        # 分推理结束和未结束两种
        if xoutput is not None:  # 结束
            xoutput = xoutput.decode('utf8')
            
            # 取本次输出
            xoutput = xoutput[xoffset:]
            xlen_output = len(xoutput)
            
            # 分输出完成未完成两种
            if xlen_output <= CHAR_LIMIT:
                # 输出完成，清理username2dynamic以标志输出完成
                rdb.hdel('username2dynamic', xuserid)
                rdb.hdel('username2offset', xuserid)
                return xoutput + PROMPT_SHORT
            else:
                # 输出未完成，分段输出
                xres = xoutput[:CHAR_LIMIT]
                
        else:  # 未结束
            
            # 从动态响应里拿
            xres = rdb.hget('username2dynamic', xuserid)
            xres = to_decode(xres)
            
            # 限制一下长度
            xres = xres[xoffset:xoffset+CHAR_LIMIT]
            
        print('output:', xres, flush=True)
        # 记录偏移
        xlen = len(xres)
        rdb.hset('username2offset', xuserid, xoffset_ + xlen)
        
        return xres + PROMPT
    
    if xinput in GO_ON_SET:  # "继续显示"
        # 从动态响应里拿
        xres = rdb.hget('username2dynamic', xuserid)
        # 没有动态响应，说明输出结束
        if xres is None:
            return do_return('（已经没有需要继续输出的内容了。）' + PROMPT_SHORT)
        
        print('GO ON >>>>', flush=True)
        users_count = rdb.scard('users_now')
        xoutput = get_result() + f'【活跃用户{users_count}人】'
        print('OUTPUT:', xoutput, flush=True)
        return xoutput
    elif xinput in RESET_SET:  # “重置对话”
        delete_many_by_user(mdb, 'dialog_in', xuserid)
        delete_many_by_user(mdb, 'dialog_out', xuserid)
        return do_return('对话已经重置。')
    
    # 生成
    rdb.sadd('users_now', xuserid)
    ts = time.time_ns()
    enqueue(mdb, 'dialog_in', xuserid, ts, {
        VALUE: xinput,
        'io': 'i',
    })
    gen = text_gen(xinput, xuserid)
    
    def generate_text():
        rdb.hdel('username2output', xuserid)
        output = ''
        for x in gen:
            rdb.hset('username2dynamic', xuserid, x)
            output = x
        rdb.hset('username2output', xuserid, output)
        rdb.srem('users_now', xuserid)
        ts = time.time_ns()
        enqueue(mdb, 'dialog_out', xuserid, ts, {
            VALUE: x,
            'io': 'o',
        })
            
    generate_text_th = threading.Thread(target=generate_text)
    generate_text_th.start()
    
    def check_quick_reply():
        """
        闭包函数，侦测快速返回，并获取。
        """
        # 最后结果
        xoutput = rdb.hget('username2output', xuserid)
        
        # 如果还没有结束
        if xoutput is None:  # 结束
            return None
        
        # 已经结束
        xoutput = xoutput.decode('utf8')
        
        # 当前偏移
        xoffset = rdb.hget('username2offset', xuserid)
        xoffset = to_int(xoffset)
        xoffset_ = xoffset
        
        # 取本次输出
        xoutput = xoutput[xoffset:]
        xlen_output = len(xoutput)
        
        # 分输出完成未完成两种
        if xlen_output <= CHAR_LIMIT:
            # 输出完成，清理username2dynamic以标志输出完成
            rdb.hdel('username2dynamic', xuserid)
            rdb.hdel('username2offset', xuserid)
            print('output:', xoutput, flush=True)
            return xoutput + PROMPT_SHORT
        else:
            # 输出未完成，分段输出
            xres = xoutput[:CHAR_LIMIT]
            print('output:', xres, flush=True)
            
            # 记录偏移
            xlen = len(xres)
            rdb.hset('username2offset', xuserid, xoffset_ + xlen)
            
            return xres + PROMPT
    
    # 既然限制5秒，干脆等3.75秒，让大模型去生成
    for _ in range(3):
        time.sleep(1.0)
        r = check_quick_reply()
        if r is not None:
            return do_return(r)
    time.sleep(0.75)
    
    # 3.75秒后，拿结果
    xres = get_result()
    return do_return(xres)


# 让服务器监听在 0.0.0.0:80
robot.config['HOST'] = '0.0.0.0'
robot.config['PORT'] = int(os.environ['W_PORT'])
robot.run()
